#include "esp_log.h"
#include "prot_port.h"
#include "version.h"
#include <string.h>
#include "driver/uart.h"
#include "driver/gpio.h"  
#include <dirent.h>
#include <sys/stat.h>
#include <inttypes.h>
#include <sys/time.h>
#include "esp_heap_caps.h"
#include "mbedtls/md5.h"

#define TAG "PROT_PORT"
#define UART_BUF_SIZE (512*1024)    // 增大UART缓冲区
#define PROTO_BUF_SIZE (512*1024)   // 增大协议缓冲区
#define UART_RX_TASK_STACK_SIZE (4 * 1024)  // 增加到4KB
#define READ_BUFFER_SIZE (1*1024*1024)  // 增加到1MB
#define FILE_BUFFER_SIZE (1*1024*1024)  // 文件读取缓冲区1MB
#define PING_PONG_BUF_SIZE (1*1024*1024)  // ping-pong buffer 大小为1MB
#define MIN(a,b) ((a) < (b) ? (a) : (b))
#define WRITE_BUFFER_SIZE (1*1024*1024)  // 1MB写入缓冲区
#define PINGPONG_BLOCK_SIZE (256*1024)  // ping-pong 缓冲区块大小
#define MD5_HASH_LENGTH 16

static uart_rx_callback g_rx_cb = NULL;
static FILE *g_write_fp = NULL;
static bool is_msc_device_connected = false;
static proto_parser_t *g_proto_parser = NULL;
static uint8_t *g_proto_buf = NULL;
static int g_uart_num = UART_NUM_1;
static protocol_handler g_proto_handler = NULL;
static FILE *g_current_fp = NULL;
static char g_current_path[256] = {0};
static uint8_t *g_read_buffer = NULL;  // 改为动态分配
static uint8_t *g_file_buffer = NULL;  // 文件读取缓冲区
static size_t g_buffer_offset = 0;
static size_t g_buffer_size = 0;
static uint8_t *g_ping_buffer = NULL;
static uint8_t *g_pong_buffer = NULL;
static SemaphoreHandle_t g_buffer_sem = NULL;
static volatile bool g_is_ping = true;
static read_data_callback g_read_cb = NULL;
static TaskHandle_t g_read_task = NULL;
static uint32_t g_read_size = 0;    // 添加读取大小变量
static uint32_t g_read_offset = 0;  // 添加读取偏移变量
// static uint8_t *g_write_buffer = NULL;  // 写入缓冲区
static mbedtls_md5_context g_md5_ctx;

void prot_set_udisk_inserted(bool set)
{
    is_msc_device_connected = set;
}

/* UART接收任务 */
static void uart_rx_task(void *arg)
{
    uint8_t *data = NULL;
    
    if (data == NULL) {
        data = (uint8_t *)malloc(UART_BUF_SIZE);
        if (!data) {
            ESP_LOGE(TAG, "Failed to allocate UART buffer");
            vTaskDelete(NULL);
            return;
        }
    }
    while (1) {
        const int rxBytes = uart_read_bytes(UART_NUM_1, data, 1, 
                                          1000 / portTICK_PERIOD_MS);
        if (rxBytes > 0) {
            // 逐字节解析
            for (int i = 0; i < rxBytes; i++) {
                if (proto_parse_byte(g_proto_parser, data[i])) {
                    // 解析成功，调用协议处理函数
                    if (g_proto_handler) {
                        g_proto_handler(g_proto_parser->cmd, 
                                      g_proto_parser->buf, 
                                      g_proto_parser->data_len);
                    }
                }
            }
        }
    }
    free(data);
    vTaskDelete(NULL);
}

bool prot_uart_send(const uint8_t *data, uint32_t len)
{
    const int txBufferSize = 8192;
    int remaining = len;
    int offset = 0;
    
    while (remaining > 0) {
        int chunk = remaining > txBufferSize ? txBufferSize : remaining;
        int sent = uart_write_bytes(g_uart_num, (const char*)data + offset, chunk);
        
        if (sent < 0) {
            ESP_LOGE(TAG, "Send failed at offset %d: %d", offset, sent);
            return false;
        }
        
        offset += sent;
        remaining -= sent;
        uart_wait_tx_done(g_uart_num, pdMS_TO_TICKS(100));
    }
    ESP_LOGI(TAG, "uart_send: offset=%d, len=%lu", offset, len);
    
    return offset == len;
}

void prot_uart_set_rx_callback(uart_rx_callback cb)
{
    g_rx_cb = cb;
}

void prot_get_device_info(struct serial_device_info *info)
{
    // 使用大端序打包版本号
    info->hw_version = (1 << 24) | (0 << 16) | (0 << 8) | 0;  // 1.0.0.0
    info->sw_version = (1 << 24) | (1 << 16) | (0 << 8) | 0;  // 1.1.0.0
    info->proto_version = (1 << 24) | (0 << 16) | (0 << 8) | 0;  // 1.0.0.0
}


bool prot_get_udisk_info(struct serial_udisk_info *info)
{
    
    return true;
}

int prot_get_file_list(const char *path, struct serial_file_info **list, int max_count)
{
    DIR *dir;
    struct dirent *ent;
    struct stat st;
    int count = 0;
    
    ESP_LOGI(TAG, "Getting file list for path: %s", path);
    
    // 预先分配固定大小的内存，避免动态增长
    *list = malloc(max_count * sizeof(struct serial_file_info));
    if (!*list) {
        ESP_LOGE(TAG, "Failed to allocate memory for file list");
        return -1;
    }
    
    if ((dir = opendir(path)) != NULL) {
        while ((ent = readdir(dir)) != NULL && count < max_count) {
            char full_path[512];
            snprintf(full_path, sizeof(full_path), "%s/%s", path, ent->d_name);
            
            if (stat(full_path, &st) == 0) {
                struct serial_file_info *info = &(*list)[count];
                strncpy(info->name, ent->d_name, sizeof(info->name) - 1);
                info->name[sizeof(info->name) - 1] = '\0';
                info->type = S_ISDIR(st.st_mode) ? 1 : 0;
                info->mtime = st.st_mtime;
                info->size = st.st_size;
                
                ESP_LOGI(TAG, "File %d: %s, type=%lu, size=%lu", 
                         count + 1, info->name, info->type, info->size);
                
                count++;
            } else {
                ESP_LOGE(TAG, "Failed to stat file: %s", full_path);
            }
        }
        closedir(dir);
        ESP_LOGI(TAG, "Found %d files/directories", count);
    } else {
        ESP_LOGE(TAG, "Failed to open directory: %s", path);
        free(*list);
        *list = NULL;
        return -1;
    }
    
    return count;
}

bool prot_mkdir(const char *path)
{
    //如果文件夹存在，则返回true
    struct stat st;
    if (stat(path, &st) == 0) {
        return true;
    }
    return mkdir(path, 0777) == 0;
}

int prot_read_file(const char *path, uint32_t offset, uint8_t *buf, uint32_t size)
{
    int read_size = 0;
    
    if (strcmp(path, g_current_path) != 0) {
        if (g_current_fp) {
            fclose(g_current_fp);
            g_current_fp = NULL;
        }
        g_current_fp = fopen(path, "rb");
        if (g_current_fp) {
            // 使用预分配的大缓冲区
            setvbuf(g_current_fp, (char*)g_file_buffer, _IOFBF, FILE_BUFFER_SIZE);
        }
        strncpy(g_current_path, path, sizeof(g_current_path)-1);
        g_buffer_size = 0;
    }
    
    if (offset >= g_buffer_offset && offset + size <= g_buffer_offset + g_buffer_size) {
        // 数据在缓冲区中，直接复制
        memcpy(buf, g_read_buffer + (offset - g_buffer_offset), size);
        read_size = size;
        
        if (offset + size > g_buffer_offset + g_buffer_size/2) {
            // 预读下一块
            fseek(g_current_fp, g_buffer_offset + g_buffer_size, SEEK_SET);
            g_buffer_size = fread(g_read_buffer, 1, READ_BUFFER_SIZE, g_current_fp);
            g_buffer_offset += g_buffer_size;
        }
    } else {
        // 数据不在缓冲区，需要读取
        if (fseek(g_current_fp, offset, SEEK_SET) != 0) {
            ESP_LOGE(TAG, "Seek failed");
            return -1;
        }
        
        read_size = fread(buf, 1, size, g_current_fp);
        
        // 预读下一块
        g_buffer_offset = offset + read_size;
        g_buffer_size = fread(g_read_buffer, 1, READ_BUFFER_SIZE, g_current_fp);
    }
    
    return read_size;
}

// 在传输结束时关闭文件
void prot_read_file_end(void)
{
    if (g_current_fp) {
        fclose(g_current_fp);
        g_current_fp = NULL;
        g_current_path[0] = '\0';
    }
}

bool prot_write_file_start(const char *path)
{
    ESP_LOGI(TAG, "Starting write file: %s", path);
    
    if (g_write_fp) {
        ESP_LOGI(TAG, "Closing previous file");
        fclose(g_write_fp);
        g_write_fp = NULL;
    }
    
    // 初始化 MD5 上下文
    mbedtls_md5_init(&g_md5_ctx);
    mbedtls_md5_starts(&g_md5_ctx);
    
    ESP_LOGI(TAG, "Opening file for writing");
    g_write_fp = fopen(path, "wb");
    if (!g_write_fp) {
        ESP_LOGE(TAG, "Failed to open file");
        return false;
    }
    
    ESP_LOGI(TAG, "File opened successfully");
    return true;
}

bool prot_write_file_data(uint32_t offset, const uint8_t *data, uint32_t size)
{
    if (!g_write_fp) {
        ESP_LOGE(TAG, "No file opened for writing");
        return false;
    }
    
    ESP_LOGI(TAG, "Writing data block: offset=%lu, size=%lu", offset, size);
    
    if (fseek(g_write_fp, offset, SEEK_SET) != 0) {
        ESP_LOGE(TAG, "Failed to seek to offset %lu", offset);
        return false;
    }
    
    size_t written = fwrite(data, 1, size, g_write_fp);
    if (written != size) {
        ESP_LOGE(TAG, "Failed to write data: written=%u, expected=%lu", written, size);
        return false;
    }
    
    // 更新 MD5
    mbedtls_md5_update(&g_md5_ctx, data, size);
    
    ESP_LOGI(TAG, "Data block written successfully");
    return true;
}

bool prot_write_file_end(uint8_t *md5_out)
{
    ESP_LOGI(TAG, "Finishing write file");
    uint8_t md5_hash[MD5_HASH_LENGTH];
    
    if (g_write_fp) {
        ESP_LOGI(TAG, "Closing file");
        fclose(g_write_fp);
        g_write_fp = NULL;
    }
    
    // 完成 MD5 计算
    mbedtls_md5_finish(&g_md5_ctx, md5_hash);
    mbedtls_md5_free(&g_md5_ctx);
    
    // 复制 MD5 值到输出缓冲区
    memcpy(md5_out, md5_hash, MD5_HASH_LENGTH);
    
    // 打印 MD5 值用于调试
    ESP_LOGI(TAG, "File MD5: ");
    for(int i = 0; i < MD5_HASH_LENGTH; i++) {
        printf("%02x", md5_hash[i]);
    }
    printf("\n");
    
    ESP_LOGI(TAG, "File closed successfully");
    return true;
}

bool prot_uart_init(int uart_num, int baud_rate)
{
    g_uart_num = uart_num;
    
    // 打印系统内存信息
    size_t free_heap = heap_caps_get_free_size(MALLOC_CAP_SPIRAM);
    size_t largest_block = heap_caps_get_largest_free_block(MALLOC_CAP_SPIRAM);
    ESP_LOGI(TAG, "System memory - Free: %d bytes, Largest block: %d bytes", 
             free_heap, largest_block);
    
    // 一次性分配所有需要的大缓冲区，并添加错误处理
    ESP_LOGI(TAG, "Allocating read buffer (%d bytes)...", READ_BUFFER_SIZE);
    g_read_buffer = heap_caps_malloc(READ_BUFFER_SIZE, MALLOC_CAP_SPIRAM);
    if (!g_read_buffer) {
        ESP_LOGE(TAG, "Failed to allocate read buffer");
        goto fail;
    }
    
    free_heap = heap_caps_get_free_size(MALLOC_CAP_SPIRAM);
    ESP_LOGI(TAG, "After read buffer - Free: %d bytes", free_heap);
    
    ESP_LOGI(TAG, "Allocating file buffer (%d bytes)...", FILE_BUFFER_SIZE);
    g_file_buffer = heap_caps_malloc(FILE_BUFFER_SIZE, MALLOC_CAP_SPIRAM);
    if (!g_file_buffer) {
        ESP_LOGE(TAG, "Failed to allocate file buffer");
        goto fail;
    }
    
    free_heap = heap_caps_get_free_size(MALLOC_CAP_SPIRAM);
    ESP_LOGI(TAG, "After file buffer - Free: %d bytes", free_heap);
    
    ESP_LOGI(TAG, "Allocating proto buffer (%d bytes)...", PROTO_BUF_SIZE);
    g_proto_buf = heap_caps_malloc(PROTO_BUF_SIZE, MALLOC_CAP_SPIRAM);
    if (!g_proto_buf) {
        ESP_LOGE(TAG, "Failed to allocate proto buffer");
        goto fail;
    }
    
    free_heap = heap_caps_get_free_size(MALLOC_CAP_SPIRAM);
    ESP_LOGI(TAG, "After proto buffer - Free: %d bytes", free_heap);
    
    ESP_LOGI(TAG, "Allocating proto parser (%d bytes)...", sizeof(proto_parser_t));
    g_proto_parser = heap_caps_malloc(sizeof(proto_parser_t), MALLOC_CAP_SPIRAM);
    if (!g_proto_parser) {
        ESP_LOGE(TAG, "Failed to allocate proto parser");
        goto fail;
    }
    
    free_heap = heap_caps_get_free_size(MALLOC_CAP_SPIRAM);
    largest_block = heap_caps_get_largest_free_block(MALLOC_CAP_SPIRAM);
    ESP_LOGI(TAG, "Final memory state - Free: %d bytes, Largest block: %d bytes",
             free_heap, largest_block);

    proto_parser_init(g_proto_parser, g_proto_buf, PROTO_BUF_SIZE);
    
    // UART配置
    uart_config_t uart_config = {
        .baud_rate = baud_rate,
        .data_bits = UART_DATA_8_BITS,
        .parity = UART_PARITY_DISABLE,
        .stop_bits = UART_STOP_BITS_1,
        .flow_ctrl = UART_HW_FLOWCTRL_DISABLE,
        .source_clk = UART_SCLK_DEFAULT
    };
    
    // 增大UART缓冲区
    ESP_ERROR_CHECK(uart_driver_install(g_uart_num, UART_BUF_SIZE * 2, 
                                      UART_BUF_SIZE * 2, 0, NULL, 0));
    ESP_ERROR_CHECK(uart_param_config(g_uart_num, &uart_config));
    
    // 设置 UART1 引脚
    uart_set_pin(UART_NUM_1, GPIO_NUM_17, GPIO_NUM_18, UART_PIN_NO_CHANGE, UART_PIN_NO_CHANGE);

    // 启用 UART DMA 模式
    uart_set_mode(g_uart_num, UART_MODE_UART);

    // 创建串口接收任务
    if (xTaskCreate(uart_rx_task, "uart_rx_task", UART_RX_TASK_STACK_SIZE,
                    NULL, 10, NULL) != pdPASS) {
        uart_driver_delete(g_uart_num);
        free(g_proto_buf);
        free(g_proto_parser);
        return false;
    }

    return true;

fail:
    // 清理已分配的内存
    if (g_read_buffer) free(g_read_buffer);
    if (g_file_buffer) {
        free(g_file_buffer);
        g_file_buffer = NULL;
    }
    if (g_proto_buf) {
        free(g_proto_buf);
        g_proto_buf = NULL;
    }
    if (g_proto_parser) {
        free(g_proto_parser);
        g_proto_parser = NULL;
    }
    return false;
}

void prot_deinit(void)
{
    // 清理 ping-pong buffer
    if (g_ping_buffer) {
        free(g_ping_buffer);
        g_ping_buffer = NULL;
    }
    if (g_pong_buffer) {
        free(g_pong_buffer);
        g_pong_buffer = NULL;
    }
    if (g_buffer_sem) {
        vSemaphoreDelete(g_buffer_sem);
        g_buffer_sem = NULL;
    }

    // 清理其他动态分配的内存
    if (g_read_buffer) {
        free(g_read_buffer);
        g_read_buffer = NULL;
    }
    if (g_file_buffer) {
        free(g_file_buffer);
        g_file_buffer = NULL;
    }
    if (g_proto_buf) {
        free(g_proto_buf);
        g_proto_buf = NULL;
    }

    // 关闭所有打开的文件
    if (g_write_fp) {
        fclose(g_write_fp);
        g_write_fp = NULL;
    }
    if (g_current_fp) {
        fclose(g_current_fp);
        g_current_fp = NULL;
    }
}

void prot_set_protocol_handler(protocol_handler handler)
{
    g_proto_handler = handler;
}

// 文件读取任务
static void read_file_task(void *arg)
{
    FILE *fp = (FILE *)arg;
    uint32_t remaining = g_read_size;
    uint32_t current_offset = g_read_offset;
    
    ESP_LOGI(TAG, "Read task started: offset=%lu, size=%lu", current_offset, remaining);
    
    while (remaining > 0) {
        // 获取当前可用的buffer
        uint8_t *current_buffer = g_is_ping ? g_ping_buffer : g_pong_buffer;
        uint32_t chunk_size = MIN(remaining, PING_PONG_BUF_SIZE);
        
        ESP_LOGI(TAG, "Reading chunk: offset=%lu, size=%lu", current_offset, chunk_size);
        
        // 读取数据
        if (fseek(fp, current_offset, SEEK_SET) != 0) {
            ESP_LOGE(TAG, "Failed to seek to offset %lu", current_offset);
            break;
        }
        
        size_t read_size = fread(current_buffer, 1, chunk_size, fp);
        if (read_size <= 0) {
            ESP_LOGE(TAG, "Failed to read data: %d", read_size);
            break;
        }
        
        ESP_LOGI(TAG, "Read %d bytes successfully", read_size);
        
        // 等待前一块数据发送完成
        xSemaphoreTake(g_buffer_sem, portMAX_DELAY);
        
        // 切换buffer
        g_is_ping = !g_is_ping;
        
        // 调用回调函数发送数据
        if (g_read_cb) {
            ESP_LOGI(TAG, "Calling read callback with %d bytes", read_size);
            g_read_cb(current_buffer, read_size);
        }
        
        remaining -= read_size;
        current_offset += read_size;
    }
    
    ESP_LOGI(TAG, "Read task completed");
    fclose(fp);
    vTaskDelete(NULL);
}

bool prot_read_file_async(const char *path, uint32_t offset, uint32_t size, read_data_callback cb)
{
    ESP_LOGI(TAG, "Starting async read: path=%s, offset=%lu, size=%lu", path, offset, size);
    
    // 检查可用内存
    uint32_t free_heap = heap_caps_get_free_size(MALLOC_CAP_SPIRAM);
    uint32_t largest_block = heap_caps_get_largest_free_block(MALLOC_CAP_SPIRAM);
    ESP_LOGI(TAG, "Memory before allocation - Free: %lu bytes, Largest block: %lu bytes", 
             free_heap, largest_block);
    ESP_LOGI(TAG, "Attempting to allocate 2 buffers of size: %d bytes each", 
             PING_PONG_BUF_SIZE);
    
    if (!g_ping_buffer) {
        ESP_LOGI(TAG, "Allocating ping-pong buffers");
        g_ping_buffer = heap_caps_malloc(PING_PONG_BUF_SIZE, MALLOC_CAP_SPIRAM);
        g_pong_buffer = heap_caps_malloc(PING_PONG_BUF_SIZE, MALLOC_CAP_SPIRAM);
        g_buffer_sem = xSemaphoreCreateBinary();
        
        if (!g_ping_buffer || !g_pong_buffer || !g_buffer_sem) {
            // 分配失败后再次检查内存
            free_heap = heap_caps_get_free_size(MALLOC_CAP_SPIRAM);
            largest_block = heap_caps_get_largest_free_block(MALLOC_CAP_SPIRAM);
            ESP_LOGE(TAG, "Memory allocation failed - Free: %lu bytes, Largest block: %lu bytes",
                     free_heap, largest_block);
            goto cleanup;
        }
    }
    
    // 打开文件
    FILE *fp = fopen(path, "rb");
    if (!fp) {
        ESP_LOGE(TAG, "Failed to open file: %s", path);
        goto cleanup;
    }
    
    ESP_LOGI(TAG, "File opened successfully");
    
    // 设置回调函数
    g_read_cb = cb;
    g_read_offset = offset;
    g_read_size = size;
    
    // 创建读取任务
    if (xTaskCreate(read_file_task, "read_task", 4096, fp, 5, &g_read_task) != pdPASS) {
        ESP_LOGE(TAG, "Failed to create read task");
        fclose(fp);
        goto cleanup;
    }
    
    ESP_LOGI(TAG, "Read task created successfully");
    
    // 释放信号量，允许第一次读取
    xSemaphoreGive(g_buffer_sem);
    return true;

cleanup:
    ESP_LOGI(TAG, "Cleaning up resources");
    if (g_ping_buffer) {
        free(g_ping_buffer);
        g_ping_buffer = NULL;
    }
    if (g_pong_buffer) {
        free(g_pong_buffer);
        g_pong_buffer = NULL;
    }
    if (g_buffer_sem) {
        vSemaphoreDelete(g_buffer_sem);
        g_buffer_sem = NULL;
    }
    return false;
}

// 数据发送完成时调用此函数
void prot_read_data_sent(void)
{
    if (g_buffer_sem) {
        xSemaphoreGive(g_buffer_sem);
    }
}
